import onnx
BATCH_SIZE = 10
onnx_model = onnx.load_model('./data/nsfw.onnx')
inputs = onnx_model.graph.input
for input in inputs:
    dim1 = input.type.tensor_type.shape.dim[0]
    dim1.dim_value = BATCH_SIZE

onnx.save_model(onnx_model, "./data/nsfw_bs{}.onnx".format(BATCH_SIZE))
print("Done saving!")